from sklearn.linear_model import LinearRegression
import joblib
import numpy as np


class Trainer(object):
    def __init__(self, dataset, model):
        (self.x_train, self.y_train, self.x_valid, self.y_valid, self.x_test) = dataset
        self.model = model.fit(self.x_train, self.y_train)

        self.y_train_p = self.model.predict(self.x_train)
        self.y_valid_p = self.model.predict(self.x_valid)
        self.y_test_p = self.model.predict(self.x_test)

        self.y_p = (self.y_train_p, self.y_valid_p, self.y_test_p)

        output = (self.y_train_p, self.y_train, self.y_valid_p, self.y_valid, self.y_test_p)
        self.output = (x.reshape(-1, 1) for x in output) if (x.ndim == 1 for x in output) else output
    #
    # def test(self, x, path):
    #     model = joblib.load(path)
    #     y_test_predictions = model.predict(x)
    #     return y_test_predictions




